"""
文本向量化模块

这个模块负责：
1. 文本向量化处理
2. 向量相似度计算
3. 批量处理优化

作者：xiaodeme@163.com
"""

import os
import logging
import numpy as np
from typing import List, Dict, Any, Optional, Tuple
from sentence_transformers import SentenceTransformer
import torch
from .config import config

# 配置日志
logging.basicConfig(level=getattr(logging, config.LOG_LEVEL))
logger = logging.getLogger(__name__)

class EmbeddingGenerator:
    """
    文本向量化生成器类
    
    负责：
    - 加载和管理向量化模型
    - 将文本转换为向量
    - 计算向量相似度
    """
    
    def __init__(self, model_name: str = "shibing624/text2vec-base-chinese"):
        """
        初始化向量化生成器
        
        Args:
            model_name (str): 使用的向量化模型名称，默认使用中文优化模型
        """
        self.model_name = model_name
        self.model = None
        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
        
        logger.info(f"初始化向量化生成器，模型: {model_name}, 设备: {self.device}")
        
        # 延迟加载模型
        self._load_model()
    
    def _load_model(self):
        """
        加载向量化模型
        """
        try:
            logger.info(f"正在加载向量化模型: {self.model_name}")
            
            # 创建模型缓存目录
            cache_dir = os.path.join(os.getcwd(), "models", "embeddings")
            os.makedirs(cache_dir, exist_ok=True)
            
            # 加载模型
            self.model = SentenceTransformer(
                self.model_name, 
                cache_folder=cache_dir,
                device=self.device
            )
            
            logger.info(f"✅ 向量化模型加载成功，向量维度: {self.model.get_sentence_embedding_dimension()}")
            
        except Exception as e:
            logger.error(f"❌ 向量化模型加载失败: {str(e)}")
            raise Exception(f"无法加载向量化模型: {str(e)}")
    
    def get_embedding_dimension(self) -> int:
        """
        获取向量维度
        
        Returns:
            int: 向量维度
        """
        if self.model is None:
            raise Exception("模型未加载")
        return self.model.get_sentence_embedding_dimension()
    
    def encode_text(self, text: str) -> np.ndarray:
        """
        将单个文本转换为向量
        
        Args:
            text (str): 输入文本
            
        Returns:
            np.ndarray: 文本向量
        """
        if not text or not text.strip():
            logger.warning("输入文本为空")
            return np.zeros(self.get_embedding_dimension())
        
        try:
            # 使用模型编码文本
            embedding = self.model.encode(text.strip(), convert_to_numpy=True)
            return embedding
            
        except Exception as e:
            logger.error(f"文本向量化失败: {str(e)}")
            raise Exception(f"文本向量化失败: {str(e)}")
    
    def encode_texts(self, texts: List[str], batch_size: int = 32, show_progress: bool = True) -> List[np.ndarray]:
        """
        批量将文本转换为向量
        
        Args:
            texts (List[str]): 文本列表
            batch_size (int): 批处理大小
            show_progress (bool): 是否显示进度
            
        Returns:
            List[np.ndarray]: 向量列表
        """
        if not texts:
            logger.warning("输入文本列表为空")
            return []
        
        logger.info(f"开始批量向量化，文本数量: {len(texts)}, 批大小: {batch_size}")
        
        try:
            # 过滤空文本
            valid_texts = [text.strip() for text in texts if text and text.strip()]
            
            if len(valid_texts) != len(texts):
                logger.warning(f"过滤了 {len(texts) - len(valid_texts)} 个空文本")
            
            # 批量编码
            embeddings = self.model.encode(
                valid_texts, 
                batch_size=batch_size,
                show_progress_bar=show_progress,
                convert_to_numpy=True
            )
            
            logger.info(f"✅ 批量向量化完成，生成了 {len(embeddings)} 个向量")
            return embeddings.tolist() if isinstance(embeddings, np.ndarray) else embeddings
            
        except Exception as e:
            logger.error(f"批量向量化失败: {str(e)}")
            raise Exception(f"批量向量化失败: {str(e)}")
    
    def calculate_similarity(self, embedding1: np.ndarray, embedding2: np.ndarray) -> float:
        """
        计算两个向量的余弦相似度
        
        Args:
            embedding1 (np.ndarray): 第一个向量
            embedding2 (np.ndarray): 第二个向量
            
        Returns:
            float: 相似度分数 (0-1之间)
        """
        try:
            # 计算余弦相似度
            similarity = np.dot(embedding1, embedding2) / (
                np.linalg.norm(embedding1) * np.linalg.norm(embedding2)
            )
            
            # 确保结果在0-1之间
            similarity = max(0.0, min(1.0, float(similarity)))
            return similarity
            
        except Exception as e:
            logger.error(f"相似度计算失败: {str(e)}")
            return 0.0
    
    def find_most_similar(self, query_embedding: np.ndarray, 
                         candidate_embeddings: List[np.ndarray], 
                         top_k: int = 5) -> List[Tuple[int, float]]:
        """
        找出最相似的向量
        
        Args:
            query_embedding (np.ndarray): 查询向量
            candidate_embeddings (List[np.ndarray]): 候选向量列表
            top_k (int): 返回最相似的K个结果
            
        Returns:
            List[Tuple[int, float]]: (索引, 相似度分数) 的列表，按相似度降序排列
        """
        if not candidate_embeddings:
            return []
        
        try:
            logger.debug(f"计算相似度，查询向量 vs {len(candidate_embeddings)} 个候选向量")
            
            similarities = []
            for i, candidate_embedding in enumerate(candidate_embeddings):
                if isinstance(candidate_embedding, list):
                    candidate_embedding = np.array(candidate_embedding)
                
                similarity = self.calculate_similarity(query_embedding, candidate_embedding)
                similarities.append((i, similarity))
            
            # 按相似度降序排序
            similarities.sort(key=lambda x: x[1], reverse=True)
            
            # 返回前K个结果
            result = similarities[:top_k]
            
            logger.debug(f"找到 {len(result)} 个最相似结果，最高相似度: {result[0][1]:.4f}")
            return result
            
        except Exception as e:
            logger.error(f"相似度计算失败: {str(e)}")
            return []

class VectorStore:
    """
    向量存储管理类
    
    负责：
    - 管理文档向量
    - 提供检索接口
    - 持久化存储
    """
    
    def __init__(self):
        """初始化向量存储"""
        self.documents = []  # 存储文档文本
        self.embeddings = []  # 存储文档向量
        self.metadata = []   # 存储文档元数据
        
        logger.info("初始化向量存储")
    
    def add_documents(self, documents: List[str], embeddings: List[np.ndarray], 
                     metadata: List[Dict[str, Any]] = None):
        """
        添加文档和对应的向量
        
        Args:
            documents (List[str]): 文档文本列表
            embeddings (List[np.ndarray]): 对应的向量列表
            metadata (List[Dict[str, Any]], optional): 文档元数据列表
        """
        if len(documents) != len(embeddings):
            raise ValueError("文档数量和向量数量不匹配")
        
        if metadata and len(metadata) != len(documents):
            raise ValueError("元数据数量和文档数量不匹配")
        
        self.documents.extend(documents)
        self.embeddings.extend(embeddings)
        
        if metadata:
            self.metadata.extend(metadata)
        else:
            self.metadata.extend([{}] * len(documents))
        
        logger.info(f"添加了 {len(documents)} 个文档，当前总计: {len(self.documents)} 个文档")
    
    def get_document_count(self) -> int:
        """
        获取文档数量
        
        Returns:
            int: 文档数量
        """
        return len(self.documents)
    
    def search(self, query_embedding: np.ndarray, top_k: int = 5, 
               similarity_threshold: float = 0.0) -> List[Dict[str, Any]]:
        """
        根据查询向量搜索相似文档
        
        Args:
            query_embedding (np.ndarray): 查询向量
            top_k (int): 返回最相似的K个文档
            similarity_threshold (float): 相似度阈值
            
        Returns:
            List[Dict[str, Any]]: 搜索结果列表
        """
        if not self.documents:
            logger.warning("向量存储中没有文档")
            return []
        
        # 使用向量化生成器的方法计算相似度
        embedding_generator = EmbeddingGenerator()
        
        similarities = []
        for i, doc_embedding in enumerate(self.embeddings):
            if isinstance(doc_embedding, list):
                doc_embedding = np.array(doc_embedding)
            
            similarity = embedding_generator.calculate_similarity(query_embedding, doc_embedding)
            
            if similarity >= similarity_threshold:
                similarities.append((i, similarity))
        
        # 按相似度降序排序
        similarities.sort(key=lambda x: x[1], reverse=True)
        
        # 构造结果
        results = []
        for i, (doc_idx, similarity) in enumerate(similarities[:top_k]):
            result = {
                'document': self.documents[doc_idx],
                'content': self.documents[doc_idx],  # 兼容性：同时提供content字段
                'similarity': similarity,
                'metadata': self.metadata[doc_idx],
                'index': doc_idx,
                'rank': i + 1
            }
            results.append(result)
        
        logger.info(f"搜索完成，找到 {len(results)} 个相关文档")
        return results
    
    def clear(self):
        """清空所有存储的数据"""
        self.documents.clear()
        self.embeddings.clear()
        self.metadata.clear()
        logger.info("向量存储已清空")

def test_embedding_generator():
    """
    测试向量化生成器
    """
    print("=== 测试向量化生成器 ===")
    
    try:
        # 创建向量化生成器
        generator = EmbeddingGenerator()
        
        # 测试单个文本向量化
        test_text = "这是一个测试文本"
        embedding = generator.encode_text(test_text)
        print(f"✅ 单文本向量化成功，向量维度: {embedding.shape}")
        
        # 测试批量文本向量化
        test_texts = ["翁氏宗谱", "家族历史", "族谱记录"]
        embeddings = generator.encode_texts(test_texts)
        print(f"✅ 批量向量化成功，生成 {len(embeddings)} 个向量")
        
        # 测试相似度计算
        emb1 = generator.encode_text("翁氏家族")
        emb2 = generator.encode_text("翁氏宗谱")
        similarity = generator.calculate_similarity(emb1, emb2)
        print(f"✅ 相似度计算成功: {similarity:.4f}")
        
        # 测试向量存储
        vector_store = VectorStore()
        vector_store.add_documents(test_texts, embeddings)
        
        # 测试搜索
        query_embedding = generator.encode_text("家族")
        results = vector_store.search(query_embedding, top_k=2)
        print(f"✅ 向量搜索成功，找到 {len(results)} 个结果")
        
        for result in results:
            print(f"   文档: {result['document']}, 相似度: {result['similarity']:.4f}")
        
        return True
        
    except Exception as e:
        print(f"❌ 测试失败: {str(e)}")
        return False

if __name__ == "__main__":
    test_embedding_generator() 